import numpy as np
import torch
from config import Config

class DataGenerator:
    def __init__(self):
        self.config = Config()
        np.random.seed(42)
        torch.manual_seed(42)
        
    def generate_features(self):
        """生成用户和物品特征"""
        user_features = torch.randn(self.config.NUM_USERS, self.config.EMBEDDING_DIM)
        item_features = torch.randn(self.config.NUM_ITEMS, self.config.EMBEDDING_DIM)
        return user_features, item_features
    
    def generate_interactions(self):
        """生成交互数据（正样本和负样本）"""
        # 正样本：每个用户随机选择10个物品
        user_ids = np.repeat(np.arange(self.config.NUM_USERS), 
                             self.config.POSITIVE_SAMPLES_PER_USER)
        pos_item_ids = np.random.randint(0, self.config.NUM_ITEMS, 
                                        size=self.config.NUM_USERS * self.config.POSITIVE_SAMPLES_PER_USER)
        pos_labels = np.ones(len(user_ids))
        
        # 负样本：每个正样本配4个随机负样本
        neg_user_ids = np.tile(user_ids, self.config.NEGATIVE_SAMPLES_PER_POSITIVE)
        neg_item_ids = np.random.randint(0, self.config.NUM_ITEMS, 
                                         size=len(neg_user_ids))
        neg_labels = np.zeros(len(neg_user_ids))
        
        # 合并数据集
        all_user_ids = np.concatenate([user_ids, neg_user_ids])
        all_item_ids = np.concatenate([pos_item_ids, neg_item_ids])
        all_labels = np.concatenate([pos_labels, neg_labels])
        
        # 打乱顺序
        indices = np.arange(len(all_user_ids))
        np.random.shuffle(indices)
        
        return (
            torch.tensor(all_user_ids[indices], dtype=torch.long),
            torch.tensor(all_item_ids[indices], dtype=torch.long),
            torch.tensor(all_labels[indices], dtype=torch.float)
        )